#include <cuda.h>
#include <omp.h>
#include <stdio.h>

#include "cush_alloc.h"
#include "cush_io.h"
#include "cuda_util.h"

/**
 *
 * Copies subsets of a buffer to the different devices.  This is done sequentially 
 * in device order.  It can be effectively thought of as a scatter operation.
 *
 * Given buffer b (size 23) and 4 devices, the uniform allocation would be:
 * b = [ * * * * * * * * * * * * * * * * * * * * * * * ]
 *      |     d0    |    d1     |    d2     |    d3   |
 *
 *  size(d0) = 6
 *  size(d1) = 6
 *  size(d2) = 6
 *  size(d3) = 5
 *
 *  Treating host_buf as a char* lets it be treated as a byte array which is all 
 *  that is needed here.
 *
 */
void cush_memcpyHtoD(void* host_buf, cush_dbuff_t* dbuff){

	int i = 0;
	//#pragma omp parallel for num_threads(dbuff->ndevices)
	for(i = 0; i < dbuff->ndevices; i++){
		cudaSetDevice(dbuff->id[i]);
        //printf("Thread %d memcpy %db to device %d from %p\n", i, dbuff->size[i], dbuff->id[i], (char*)host_buf+dbuff->offset[i]);
		cudaMemcpyAsync(dbuff->ptr[i], 
					(char*)host_buf+dbuff->offset[i], 
					dbuff->size[i], 
					cudaMemcpyHostToDevice,
                    dbuff->stream[i]);
		//cuMemcpyHtoD(dbuff->ptr[i], host_buf+offset, dbuff->size[i]);
		checkCUDAError("cush memcpyHtoD failed", i);
	}
}

void cush_memcpyHtoD(void* host_buf, cush_dbuff_t* dbuff, int block_size, int stride){

	int i = 0;
	//#pragma omp parallel for num_threads(dbuff->ndevices)
	for(i = 0; i < dbuff->ndevices; i++){
		//int th_id = omp_get_thread_num();
		cudaSetDevice(dbuff->id[i]);
        //printf("Thread %d memcpy to devices %d\n", th_id, dbuff->id[th_id]);
		cudaMemcpyAsync(dbuff->ptr[i], 
					(char*)host_buf+dbuff->offset[i], 
					dbuff->size[i], 
					cudaMemcpyHostToDevice,
                    dbuff->stream[i]);
		//cuMemcpyHtoD(dbuff->ptr[i], host_buf+offset, dbuff->size[i]);
		checkCUDAError("cush memcpy failed", i);
	}
}

/**
 *
 * Copies subsets of a buffer that reside on different devices back to a single buffer 
 * on the host.  This is done sequentially in device order.  It can be effectively 
 * thought of as a gather operation.
 *
 * For example, this function would gather data d0-d3 into buffer b (size 23) from 4 devices:
 * b = [ * * * * * * * * * * * * * * * * * * * * * * * ]
 *      |     d0    |    d1     |    d2     |    d3   |
 *
 *  size(d0) = 6
 *  size(d1) = 6
 *  size(d2) = 6
 *  size(d3) = 5
 *
 *  Treating host_buf as a char* lets it be treated as a byte array which is all 
 *  that is needed here.
 */
void cush_memcpyDtoH(void* host_buf, cush_dbuff_t* dbuff){

	int i = 0;
	//#pragma omp parallel for num_threads(dbuff->ndevices)
	for(i = 0; i < dbuff->ndevices; i++){
		//int th_id = omp_get_thread_num();
		cudaSetDevice(dbuff->id[i]);

        //printf("\ndbuf = %lu\n", dbuff->ptr[th_id]);
        //printf("hbuf = %lu\n", host_buf);
        //printf("offset = %d\n", dbuff->offset[th_id]);
        //printf("args: %lu,\t%lu,\t%d,\tDtoH\n", (char*)host_buf+dbuff->offset[th_id], dbuff->ptr[th_id], dbuff->size[th_id]);
        //printf("copying from %lu to %lu\n",  (char*)dbuff->ptr[th_id]+dbuff->offset[th_id],(void**)((char*)host_buf+dbuff->offset[th_id]));
		cudaMemcpyAsync(((char*)host_buf+dbuff->offset[i]), 
					dbuff->ptr[i], 
					dbuff->size[i], 
					cudaMemcpyDeviceToHost,
                    dbuff->stream[i]);
		checkCUDAError("cush memcpy failed", i);
	}
}

void cush_memcpyDtoD(cush_dbuff_t* dst_buff, cush_dbuff_t* src_buff){

    if(src_buff->ndevices != dst_buff->ndevices){
        fprintf(stderr, "Unequal number of devices involved in cush_all-to-all\n");
        exit(EXIT_FAILURE);
    }

    int src_id = src_buff->id[0];
    int dst_id = dst_buff->id[0];

    int msg_size = src_buff->size[0];

    cudaSetDevice(dst_id);
    checkCUDAError("Setting all-to-all device failed", src_id);

    cudaMemcpyPeerAsync( (void*)((char*)(dst_buff->ptr[0])),
                                dst_id, 
                                (void*)((char*)(src_buff->ptr[0])),
                                src_id, 
                                msg_size);
}


/* vim: set sw=4 sts=4 et foldmethod=syntax syntax=c : */

